import multiprocessing
from concurrent.futures import ProcessPoolExecutor
from typing import Set
import hashlib
import addresses
import sys
import os
import numpy as np
import time
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(current_dir, "lib"))
import ecdsa
import base58

# 导入CUDA相关库
try:
    import pycuda.driver as cuda
    import pycuda.autoinit
    from pycuda.compiler import SourceModule
    CUDA_AVAILABLE = True
except ImportError:
    print("CUDA libraries not available. Please install pycuda.")
    CUDA_AVAILABLE = False

TARGET_ADDRESSES: Set[str] = set(addresses.TARGET_ADDRESSES)
CPU_HALF = False
MIN_WORKERS = 1
BATCH_SIZE = 10000  # 增大批次以充分利用GPU并行能力
USE_GPU = True  # 控制是否使用GPU
TARGET_HASHES = {hashlib.sha256(hashlib.sha256(base58.b58decode(address)).digest()).digest()[:4] 
                for address in TARGET_ADDRESSES}

# CUDA核函数代码
if CUDA_AVAILABLE:
    cuda_code = """
    #include <stdio.h>
    #include <stdint.h>
    
    // 简化版随机数生成（实际应使用更好的随机数生成器）
    __device__ void generate_random(unsigned char *output, int length, uint64_t seed) {
        uint64_t x = seed;
        for (int i = 0; i < length; i++) {
            x = x * 6364136223846793005ULL + 1442695040888963407ULL;
            output[i] = (unsigned char)(x >> 32);
        }
    }
    
    // 简化版SHA256哈希函数（实际实现需要完整SHA256算法）
    __device__ void sha256_simple(const unsigned char *input, int input_len, unsigned char *output) {
        // 这只是一个占位符，实际应实现完整的SHA256
        for (int i = 0; i < 32; i++) {
            output[i] = input[i % input_len] ^ (input[(i+1) % input_len] + i);
        }
    }
    
    __global__ void generate_keys(
        unsigned char *private_keys,
        unsigned char *public_keys,
        unsigned char *addresses,
        int *found_flags,
        uint64_t base_seed,
        int batch_size
    ) {
        int idx = threadIdx.x + blockIdx.x * blockDim.x;
        if (idx < batch_size) {
            // 生成私钥
            unsigned char private_key[32];
            generate_random(private_key, 32, base_seed + idx);
            
            // 复制私钥到输出数组
            for (int i = 0; i < 32; i++) {
                private_keys[idx * 32 + i] = private_key[i];
            }
            
            // 模拟生成公钥（实际应使用椭圆曲线算法）
            unsigned char public_key[65];
            public_key[0] = 0x04;  // 未压缩公钥前缀
            sha256_simple(private_key, 32, public_key + 1);
            sha256_simple(public_key + 1, 32, public_key + 33);
            
            // 复制公钥到输出数组
            for (int i = 0; i < 65; i++) {
                public_keys[idx * 65 + i] = public_key[i];
            }
            
            // 模拟生成地址（实际应使用RIPEMD160和Base58）
            unsigned char address_bytes[25];
            sha256_simple(public_key, 65, address_bytes);
            
            // 复制地址到输出数组
            for (int i = 0; i < 25; i++) {
                addresses[idx * 25 + i] = address_bytes[i];
            }
            
            // 模拟检查地址是否匹配（实际应与TARGET_HASHES比较）
            // 这里只是示例，实际逻辑需要更复杂
            if ((address_bytes[0] & 0x0F) == 0) {  // 简单示例：如果首字节低4位为0则视为匹配
                found_flags[idx] = 1;
            } else {
                found_flags[idx] = 0;
            }
        }
    }
    """
    
    # 编译CUDA核函数
    try:
        cuda_module = SourceModule(cuda_code)
        generate_keys_kernel = cuda_module.get_function("generate_keys")
    except Exception as e:
        print(f"CUDA compilation error: {e}")
        CUDA_AVAILABLE = False

def generate_key_pair():
    private_key=os.urandom(32)
    sk = ecdsa.SigningKey.from_string(private_key, curve=ecdsa.SECP256k1)
    vk = sk.get_verifying_key()
    public_key = b'\x04' + vk.to_string()
    ripemd160 = hashlib.new('ripemd160')
    ripemd160.update(hashlib.sha256(public_key).digest())
    address_hash = b'\x00' + ripemd160.digest()
    checksum = hashlib.sha256(hashlib.sha256(address_hash).digest()).digest()[:4]
    address = base58.b58encode(address_hash + checksum).decode('utf-8')
    return private_key.hex(), public_key.hex(), address

def generate_key_pairs_batch():
    if USE_GPU and CUDA_AVAILABLE:
        # 使用GPU生成密钥对
        batch_size = BATCH_SIZE
        
        # 分配GPU内存
        private_keys = np.zeros(batch_size * 32, dtype=np.uint8)
        public_keys = np.zeros(batch_size * 65, dtype=np.uint8)  # 未压缩公钥长度为65字节
        addresses = np.zeros(batch_size * 25, dtype=np.uint8)    # 地址长度假设为25字节
        found_flags = np.zeros(batch_size, dtype=np.int32)
        
        # 生成随机种子
        base_seed = np.uint64(int(time.time() * 1000000))
        
        # 调用CUDA核函数
        block_size = 256
        grid_size = (batch_size + block_size - 1) // block_size
        
        generate_keys_kernel(
            cuda.InOut(private_keys),
            cuda.InOut(public_keys),
            cuda.InOut(addresses),
            cuda.InOut(found_flags),
            np.uint64(base_seed),
            np.int32(batch_size),
            block=(block_size, 1, 1),
            grid=(grid_size, 1)
        )
        
        # 处理结果
        key_pairs = []
        for i in range(batch_size):
            if found_flags[i] == 1:  # 只处理可能匹配的结果
                priv_key = private_keys[i*32:(i+1)*32].tobytes().hex()
                pub_key = public_keys[i*65:(i+1)*65].tobytes().hex()
                # 注意：这里简化处理，实际应该将地址字节转换为Base58格式
                addr_bytes = addresses[i*25:(i+1)*25].tobytes()
                # 临时使用简化的Base58编码（实际应使用完整的Base58Check）
                addr = base58.b58encode(addr_bytes).decode('utf-8')
                key_pairs.append((priv_key, pub_key, addr))
        
        # 如果没有找到匹配项，至少返回一些结果用于检查
        if not key_pairs and batch_size > 0:
            # 返回前10个结果用于验证
            for i in range(min(10, batch_size)):
                priv_key = private_keys[i*32:(i+1)*32].tobytes().hex()
                pub_key = public_keys[i*65:(i+1)*65].tobytes().hex()
                addr_bytes = addresses[i*25:(i+1)*25].tobytes()
                addr = base58.b58encode(addr_bytes).decode('utf-8')
                key_pairs.append((priv_key, pub_key, addr))
        
        return key_pairs
    else:
        # 使用CPU生成密钥对（原始方法）
        return [generate_key_pair() for _ in range(BATCH_SIZE)]

def check_and_save_batch(key_pairs):
    found = False
    with open("checklist.txt", "a") as f:
        for private_key, public_key, address in key_pairs:
            try:
                # 对于GPU生成的结果，可能需要额外验证
                address_bytes = base58.b58decode(address)
                address_hash = hashlib.sha256(hashlib.sha256(address_bytes).digest()).digest()[:4]
                
                if address_hash in TARGET_HASHES and address in TARGET_ADDRESSES:
                    f.write(f"{private_key}-{public_key}-{address}\n")
                    print(f"Found match: {address}")
                    found = True
            except Exception as e:
                # 处理可能的解码错误
                print(f"Error checking address {address}: {e}")
                continue
    return found

def worker():
    while True:
        key_pairs = generate_key_pairs_batch()
        if check_and_save_batch(key_pairs):
            print(f"Found matching address in batch")
            break

def main():
    print("Starting")
    
    # 检查是否使用GPU
    if USE_GPU and CUDA_AVAILABLE:
        print("GPU mode enabled - using CUDA for key generation")
        # GPU模式下，只需要一个工作进程，因为GPU已经在内部并行处理
        num_workers = 1
    else:
        print("CPU mode - GPU not available or disabled")
        num_cores = multiprocessing.cpu_count()
        if CPU_HALF:
            num_workers = max(int(num_cores / 2), MIN_WORKERS)
            print("Low CPU mode")
        else:
            num_workers = max(num_cores - 1, MIN_WORKERS)
            print("High CPU mode")
    
    print(f"Using {num_workers} worker processes")
    print(f"Target addresses: {len(TARGET_ADDRESSES)}")
    print(f"Batch size: {BATCH_SIZE}")
    
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(worker) for _ in range(num_workers)]
        for future in futures:
            future.result()

if __name__ == "__main__":
    main()